import datetime
import json
import pandas as pd
import random
from collections import defaultdict
from fastapi import APIRouter, Depends
from sqlalchemy.orm import Session

from database.database import get_db
from models.custom_history import CustomHistory
from schemas.gun_param import ParamRequires
from schemas.user import User

router = APIRouter()

@router.post("/", tags=["alg"])
async def genetic_algorithm(user: User, gun_param: ParamRequires, population_size=100, generations=100, mutation_rate=0.1, db: Session=Depends(get_db)):
    print("接收到的数据:", gun_param.model_dump())
    # 获取数据
    excel_file = pd.ExcelFile(r"C:\Users\Administrator\Documents\改枪.xlsx")
    data = pd.read_excel(excel_file)
    param_weights = gun_param.param_weights.model_dump()
    add_values = gun_param.add_values.model_dump()
    price_range = gun_param.price_range.model_dump()
    must_have_part = set(gun_param.must_have_part)
    must_not_have_part = set(gun_param.must_not_have_part)
    must_have_accessory = set(gun_param.must_have_accessory)
    must_not_have_accessory = set(gun_param.must_not_have_accessory)
    
    ## 数据处理
    # 性能权重处理
    for param, weight in param_weights.items():
        if weight is None:
            param_weights[param] = 0
    # 如果某个性能没有提供权重，但是提供了要求，那么将该要求置为None
    for param, performance in add_values.items():
        if performance is not None and performance > 0 and param_weights[param] == 0:
            add_values[param] = None
    # 如果没有提供性能权重，也没有提供性能要求，那么性能权重平均分配
    if all(value == 0 for value in param_weights.values()) and all(value is None for value in add_values.values()):
        equal_weight = 1 / len(param_weights)
        param_weights = {performance: equal_weight for performance in param_weights.keys()}
    # 如果提供了性能要求，但是之和不为1，那么进行归一化
    if sum(param_weights.values()):
        total_weight = sum(param_weights.values())
        param_weights = {performance: weight / total_weight for performance, weight in param_weights.items()}
    # 处理价格范围
    if price_range['min_'] is None:
        price_range['min_'] = 0
    if price_range['max_'] is None:
        price_range['max_'] = 100000000
    price_range = {key: value * 10000 for key, value in price_range.items()}
    min_price = price_range['min_']
    max_price = price_range['max_']
    # 处理性能要求
    for performance, value in add_values.items():
        if value is None:
            add_values[performance] = -100.0

    # 将数据框转换为字典，方便快速访问
    accessory_info = {}
    for _, row in data.iterrows():
        accessory = row['accessory']
        accessory_info[accessory] = {
            'price': row['price'],
            'performance': {
                performance: row[performance]
                for performance in add_values.keys()
                if pd.notnull(row[performance])
            },
            'part': row['part']
        }

    # 获取所有部位
    parts = data['part'].unique()
    part_available_accessories = {
        part: data[data['part'] == part]['accessory'].tolist()
        for part in parts
    }

    # 创建初始种群
    def create_individual():
        individual = {}
        for part in parts:
            available_accessories = part_available_accessories[part]
            if part in must_have_part:
                individual[part] = random.choice(available_accessories)
            elif part in must_not_have_part:
                individual[part] = None
            else:
                individual[part] = random.choice(available_accessories + [None])
        for accessory in must_have_accessory:
            part = accessory_info[accessory]['part']
            individual[part] = accessory
        for accessory in must_not_have_accessory:
            part = accessory_info[accessory]['part']
            individual[part] = None
        return individual

    population = [create_individual() for _ in range(population_size)]

    # 评估函数
    def evaluate(individual):
        total_price = 0
        performance_values = defaultdict(int)
        # 计算总价格和性能值
        for accessory in individual.values():
            if accessory:
                info = accessory_info[accessory]
                total_price += info['price']
                # 直接更新性能值，减少循环嵌套
                performance_values.update(
                    (k, performance_values[k] + v) for k, v in info['performance'].items()
                )

        # 检查价格要求
        if total_price < min_price or total_price > max_price:
            return 0

        # 检查性能约束
        for performance, constraint in add_values.items():
            if performance_values[performance] < constraint:
                return 0

        # 检查配件要求
        for accessory in must_have_accessory:
            if individual[accessory] is None:
                return 0
        for accessory in must_not_have_accessory:
            if individual[accessory] is not None:
                return 0

        # 检查部位要求
        for part in must_have_part:
            if individual[part] is None:
                return 0
        for part in must_not_have_part:
            if individual[part] is not None:
                return 0

        # 检查增高座瞄具和瞄具的关系
        if individual['增高座瞄具'] is not None and individual['瞄准镜'] is None:
            return 0

        # 检查枪管和贴片的关系
        if individual['枪管'] is None:
            for patch in ['左贴片', '右贴片', '上贴片']:
                if individual[patch] is not None:
                    return 0

        fitness = sum(
            param_weights[performance] * performance_values[performance]
            for performance in param_weights.keys()
            if performance in performance_values
        )
        return fitness

    # 选择函数（轮盘赌选择）
    def select(population):
        fitness_scores = [evaluate(individual) for individual in population]
        total_fitness = sum(fitness_scores)
        if total_fitness == 0:
            return random.choices(population, k=2)
        probabilities = [score / total_fitness for score in fitness_scores]
        return random.choices(population, weights=probabilities, k=2)

    # 交叉函数
    def crossover(parent1, parent2):
        child = {}
        parts = list(parent1.keys())
        crossover_point = random.randint(1, len(parts) - 1)
        for i in range(crossover_point):
            child[parts[i]] = parent1[parts[i]]
        for i in range(crossover_point, len(parts)):
            child[parts[i]] = parent2[parts[i]]
        return child

    # 变异函数
    def mutate(individual):
        for part in individual.keys():
            if random.random() < mutation_rate:
                available_accessories = part_available_accessories[part]
                if part in must_have_part:
                    individual[part] = random.choice(available_accessories)
                elif part in must_not_have_part:
                    individual[part] = None
                else:
                    individual[part] = random.choice(available_accessories + [None])
            return individual

    # 遗传算法主循环
    for i in range(generations):
        print(i)
        new_population = []
        for _ in range(population_size // 2):
            parent1, parent2 = select(population)
            child1 = crossover(parent1, parent2)
            child2 = crossover(parent2, parent1)
            child1 = mutate(child1)
            child2 = mutate(child2)
            new_population.extend([child1, child2])
        population = new_population

    best_individual = max(population, key=evaluate)
    best_fitness = evaluate(best_individual)

    condition = gun_param.model_dump_json()
    if best_fitness == 0:
        print("不存在符合要求的配件方案")
        custom_history = CustomHistory(
            open_id=user.open_id,
            result = json.dumps({}, ensure_ascii=False),
            condition = condition,
            create_time = datetime.datetime.now()
        )
        db.add(custom_history)
        db.commit()
        db.refresh(custom_history)
        return None
    else:
        print("最优配件方案: ", best_individual)
        custom_history = CustomHistory(
            open_id=user.open_id,
            result = json.dumps(best_individual, ensure_ascii=False),
            condition = condition,
            create_time = datetime.datetime.now()
        )
        db.add(custom_history)
        db.commit()
        db.refresh(custom_history)
        return best_individual
    

if __name__ == "__main__":
    # 各项性能权重，可为0
    weight_coefficients = {
        # '优势射程': 0.1,
        # '后坐力控制': 0.4,
        # '操控速度': 0.2,
        # '据枪稳定性': 0.2,
        # '腰际射击精度': 0.1
    }
    # 性能要求，可不填
    performance_constraints = {
        '后坐力控制': 10
    }
    # 价格要求，可不填
    price_requirement = (100000, 200000)
    # 部位要求，可不填
    part_requirements = {
        '枪口':'must'
    }
    # 配件要求，可不填
    accessory_requirements = {
        # '薪火螺旋消焰器':'must',
        # '先进多口径消音器':'must_not'
    }

    # 运行遗传算法
    import asyncio
    result = asyncio.run(genetic_algorithm(weight_coefficients, performance_constraints, price_requirement, part_requirements, accessory_requirements))

    if result:
        print("最优配件方案:")
        for part, accessory in result.items():
            print(f"{part}: {accessory}")
    else:
        print("不存在符合要求的配件方案")